import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision import transforms
from torchvision.transforms import ToTensor
import torchvision.transforms as tt
import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn import metrics

import os

# 导入自己创建的python文件
import sys
sys.path.append("..") # Adds higher directory to python modules path.
from frame.ModelUtil import *
from frame.DataProcess import *
from frame.TrainUtil import train, train_models

LEARNING_RATE = 5e-4
BATCH_SIZE = 128
MODEL = 'NN_4layer'
EPOCHS = 100
DATA_NAME = 'Purchase100_limited'
num_shadowsets = 100
seed = 0
prop_keep = 0.5

# 批量训练模型
weight_dir = os.path.join('..', 'weights_for_exp', DATA_NAME)
# transform = transforms.Compose([])
transform = transforms.Compose([])
# transform = transforms.Compose([
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.5], std=[0.5])
#     ])
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")


num_shadowsets_list = [2,4,6,8,10,20,30,40,50]
for num in num_shadowsets_list:
    train_models(weight_dir, DATA_NAME, MODEL, LEARNING_RATE, EPOCHS, transform, device, num_shadowsets=num, prop_keep=0.5)